import torchvision.transforms as T
import torchvision
import torch
from Models import densenet
from Models import resnet
from Models import vgg
from Util import  Savetar
import torch.nn as nn

def predict_image(model, image):
    output = model(image)
    output = torch.argmax(output, dim=1).item()
    return output

map_location = None
densenet121 = densenet.densenet121(pretrained=False)
resnet50 = resnet.resnet50(pretrained=False)
vgg19bn = vgg.vgg19_bn(pretrained=False)
densenet121.eval()
resnet50.eval()
vgg19bn.eval()


state_dicts_densenet121 = torch.load("./Models/state_dicts/densenet121.pth",map_location=map_location)
densenet121.load_state_dict(state_dicts_densenet121['net'])
densenet121.eval()

state_dicts_resnet50 = torch.load("./Models/state_dicts/resnet50.pth",map_location=map_location)
resnet50.load_state_dict(state_dicts_resnet50['net'])
resnet50.eval()

state_dicts_vggg19bn = torch.load("./Models/state_dicts/vgg19.pth",map_location=map_location)
vgg19bn.load_state_dict(state_dicts_vggg19bn['net'])
vgg19bn.eval()


transform = T.Compose([T.ToTensor(),T.Normalize(mean = [0.4914, 0.4822, 0.4465], std = [0.2471, 0.2435, 0.2616])])#
testset = torchvision.datasets.CIFAR10(root='./data', train=False,download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=1, shuffle=True, num_workers=0)

def get_lib():
    data = dict(zip([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [[]] * 10))

    for i,(images,label) in enumerate(testloader):

        flag = 0

        for i in range(10):
            if len(data[i]) == 100:
                flag += 1
        if flag == 10:
            break

        if predict_image(densenet121,images)!=label.item():
            continue
        if predict_image(resnet50,images)!=label.item():
            continue
        if predict_image(vgg19bn,images)!=label.item():
            continue
        if len(data[label.item()]) == 100:
            continue


        temp = data[label.item()].copy()
        temp.append(images.detach().numpy().tolist())
        data.update({label.item():temp})

        Savetar(data={label.item(): data[label.item()]},  # 代保存字典文件
                path_file='./library/library1.json',
                type='write'
                )
    return data

if __name__ == "__main__":
    get_lib()




